# Copyright (c) 2020 cgsdfc
#
# This software is released under the MIT License.
# https://opensource.org/licenses/MIT

from collections import namedtuple
import enum
from typing import Sequence, Union
import abc


def HackerCaseToCamelCase(s: str):
    return ''.join(map(str.capitalize, s.split('_')))


NameBundle = namedtuple('NameBundle', ['asdl_name', 'cc_name', 'c_name'])


class Node:
    def __init__(self, name: NameBundle):
        self._name = name

    @property
    def name(self):
        return self._name.cc_name

    @property
    def asdl_name(self):
        return self._name.asdl_name

    @property
    def c_name(self):
        return self._name.c_name


class ValueSemantic(enum.Enum):
    BY_VALUE = 0
    BY_PTR = 1
    BY_CONST_REF = 2


class NumArgs(enum.Enum):
    REQUIRED = 0
    OPTIONAL = 1
    SEQUENCE = 2


class Member(Node):
    @classmethod
    def MangleGetterName(cls, asdl_name):
        "`format_spec` => `FormatSpec`"
        return HackerCaseToCamelCase(asdl_name)

    @classmethod
    def MangleSetterName(cls, asdl_name):
        "`format_spec` => `SetFormatSpec`"
        return 'Set' + HackerCaseToCamelCase(asdl_name)

    @classmethod
    def MangleName(cls, asdl_name):
        # currently there is nothing special.
        return NameBundle(asdl_name, asdl_name, asdl_name)

    def __init__(
            self,
            type,
            asdl_name,
            nargs=NumArgs.REQUIRED,
            # TODO: this should go into types.
            value_semantic=ValueSemantic.BY_VALUE):

        super().__init__(self.MangleName(asdl_name))
        self.type = type
        self._nargs = nargs
        self.value_semantic = value_semantic

    def CMemberName(self):
        return self._name.c_name

    def TypeName(self):
        return self.type.TypeName()

    def TypeNameWithValueSemantic(self):
        if self.value_semantic == ValueSemantic.BY_PTR:
            return Ptr(self.typename)
        if self.value_semantic == ValueSemantic.BY_VALUE:
            return self.typename
        if self.value_semantic == ValueSemantic.BY_CONST_REF:
            return ConstRef(self.typename)

    def FormalArg(self):
        return self.TypeNameWithValueSemantic() + ' ' + self.name

    def ActualArg(self):
        return ConvertToC(self.type, self.name)

    @property
    def typename(self):
        return self.type.name

    @property
    def getter_name(self):
        return self.MangleGetterName(self._name.asdl_name)

    @property
    def setter_name(self):
        return self.MangleSetterName(self._name.asdl_name)

    @property
    def is_required(self):
        return self._nargs == NumArgs.REQUIRED

    @property
    def is_optional(self):
        return self._nargs == NumArgs.OPTIONAL

    @property
    def is_sequence(self):
        return self._nargs == NumArgs.SEQUENCE


class AstNode(Node):
    # To be compatible with pybind.
    GET_UNTYPED_RAW_PTR = 'GetUntypedRawPtr'

    GET_RAW_PTR = 'ptr'

    GET_ARENA = 'GetArena'

    def __init__(self, name: NameBundle, parent: 'AstNode'):
        super().__init__(name)
        self._parent = parent

    def TypeName(self):
        "Class name for cc code"
        return self._name.cc_name

    def New(self, raw_ptr, arena=None):
        if arena is None:
            arena = self.GetArena()
        return '{type}({raw_ptr}, {arena})'.format(type=self.name,
                                                   raw_ptr=raw_ptr,
                                                   arena=arena)

    @classmethod
    def GetRawPtr(self):
        return self.GET_RAW_PTR + '()'

    @classmethod
    def GetUntypedRawPtr(self):
        return self.GET_UNTYPED_RAW_PTR + '()'

    @classmethod
    def GetArena(self):
        return self.GET_ARENA + '()'

    @property
    def parent(self):
        return self._parent.name

    def UnderlyingType(self):
        return self._name.c_name

    @property
    def underlying_type(self):
        return self.UnderlyingType()


class AncestorAstNode(AstNode):
    "Represents the top of the hierarchy -- AstNode class."
    TYPENAME = 'AstNode'

    def __init__(self):
        super().__init__(NameBundle(None, self.TYPENAME, None), None)


# The only instance of `AncestorAstNode`
AST_NODE = AncestorAstNode()


def Ptr(typename):
    return typename + '*'


def ConstRef(typename):
    return 'const {}&'.format(typename)


def StaticCast(tparam, param):
    return 'static_cast<{tparam}>({param})'.format(tparam=tparam, param=param)


def C_Cast(tparam, param):
    return '({tparam}){param}'.format(tparam=tparam, param=param)


def ReinterpretBorrow(tparam, param):
    return 'pybind11::reinterpret_borrow<{tparam}>({param})'.format(
        tparam=tparam, param=param)


def Construct(type, *args):
    return type + '{' + ', '.join(args) + '}'


def ConvertToCXX(type, ref):
    "Given a `ref`, convert it to a CXX expression of `type`."

    if isinstance(type, Object):
        return ReinterpretBorrow(type.name, ref)
    if isinstance(type, (AstNode, AsdlSeq)):
        return Construct(type.name, ref, type.GetArena())
    if isinstance(type, Enum):
        return StaticCast(type.name, ref)
    return ref


def ConvertToC(type, ref):
    if isinstance(type, Enum):
        return StaticCast(type.c_name, ref)
    if isinstance(type, (AstNode, Object)):
        return '{ref}.{method}'.format(ref=ref, method=type.GetRawPtr())
    return ref


class FinalAstNode(AstNode):
    "A node that is final -- holding concrete data (members)."

    def __init__(self, name: NameBundle, parent: AstNode,
                 members: Sequence[Member]):
        super().__init__(name, parent)
        self.members = members

    def PyCreateFunction(self):
        "The CPython function that create this node."
        return '_Py_{0.asdl_name}'.format(self)

    def ConstructorFormalArgs(self):
        "Formal arguments list of constructor."
        return ', '.join(m.FormalArg() for m in self.members)

    def ConstructorActualArgs(self):
        return ', '.join(m.ActualArg() for m in self.members)

    @abc.abstractmethod
    def MemberRef(self, member: Member):
        pass

    def MemberGetter(self, member: Member):
        """
        Return a converted getter expression, like:
        Member `value` on `ReturnNode` => `ReturnNode{ptr()->v.Return.value, GetArena()}`
        """
        return ConvertToCXX(member.type, self.MemberRef(member))

    def MemberSetter(self, member: Member, new_value: str):
        """
        Return a converted setter expression, like:
        Member `value` on `ReturnNode` yields:

        ptr()->v.Return.value = value.ptr()
        """
        return '{ref} = {value}'.format(ref=self.MemberRef(member),
                                        value=ConvertToC(
                                            member.type, new_value))

    @abc.abstractmethod
    def HasLocation(self):
        pass

    @abc.abstractmethod
    def HasEndLocation(self):
        pass

    def LocationArgsCount(self):
        return 0 if not self.HasLocation() else (
            2 if not self.HasEndLocation() else 4)

    def CallPyCreateFunc(self, arena: str):
        func = self.PyCreateFunction()
        location = ['0'] * self.LocationArgsCount()
        actual_args = ', '.join([m.ActualArg() for m in self.members] +
                                location + [arena + '->ptr()'])
        return '{func}({args})'.format(func=func, args=actual_args)

    @property
    def py_create_func(self):
        return self.PyCreateFunction()

    @property
    def formal_args(self):
        return self.ConstructorFormalArgs()

    @property
    def actual_args(self):
        return self.ConstructorActualArgs()

    def MaybeExplicit(self):
        "If the constructor needs an `explicit`, return that."
        return 'explicit' if 0 == len(self.members) else ''

    @property
    def explicit_specifier(self):
        return self.MaybeExplicit()


class Location:
    "Class dedicated to handle `Location`."
    LOCATION_TYPENAME = 'Location'

    LOCATION_LINENO = 'lineno'
    LOCATION_COL_OFFSET = 'col_offset'
    LOCATION_FIELDS = (LOCATION_LINENO, LOCATION_COL_OFFSET)

    LOCATION_END_LINENO = 'end_lineno'
    LOCATION_END_COL_OFFSET = 'end_col_offset'

    LOCATION_END_FIELDS = (LOCATION_END_LINENO, LOCATION_END_COL_OFFSET)

    def GetRawPtr(self):
        raise NotImplemented('Derived class must implement this method.')

    def LocationFieldRef(self, field: str):
        return '{ptr}->{field}'.format(ptr=self.GetRawPtr(), field=field)

    def LocationGetter(self, fields: Sequence[str]):
        return '{type}({args})'.format(type=Location.LOCATION_TYPENAME,
                                       args=', '.join(
                                           map(self.LocationFieldRef, fields)))

    def LocationSetter(self, fields, new_value):

        def SetField(lhs_field, rhs_field):
            return '{ref} = {value}.{field};'.format(
                ref=self.LocationFieldRef(lhs_field),
                value=new_value,
                field=rhs_field)

        return '\n'.join([SetField(lhs, rhs) for lhs, rhs in zip(fields, self.LOCATION_FIELDS)])

    def __init__(self, has_location: bool, has_end_location: bool):
        self.has_location = has_location
        self.has_end_location = has_end_location


class TypeNode(AstNode, Location):

    GET_KIND = 'GetKind'
    ENUM_TYPE = 'Kind'

    @classmethod
    def MangleName(cls, asdl_name):
        return NameBundle(asdl_name,
                          cc_name=asdl_name + '_node',
                          c_name=asdl_name + '_ty')

    def __init__(self, asdl_name, kind_enum: 'KindEnum', has_location: bool,
                 has_end_location: bool):
        assert isinstance(kind_enum, KindEnum)
        AstNode.__init__(self,
                         name=self.MangleName(asdl_name),
                         parent=AST_NODE)
        Location.__init__(self, has_location, has_end_location)
        self._kind_enum = kind_enum

    @classmethod
    def GetKind(self):
        return 'GetKind()'

    @property
    def enum_type(self):
        "For `expr_node`, this will return `expr_kind`"
        return self._kind_enum.name


class ConstructorNode(FinalAstNode):
    """
    Things like:
    `FormattedValue(expr value, int? conversion, expr? format_spec)`
    """

    # The variable name used in the CPython's code as the union.
    UNION_NAME = 'v'

    # The name of the enum value variable that denotes this class.
    ENUM_VALUE = 'kEnumValue'

    @classmethod
    def MangleName(cls, asdl_name):
        assert asdl_name[0].isupper()
        return NameBundle(asdl_name,
                          cc_name=asdl_name + 'Node',
                          c_name=asdl_name)

    def __init__(self, asdl_name, parent: TypeNode, members: Sequence[Member]):
        super().__init__(name=self.MangleName(asdl_name),
                         parent=parent,
                         members=members)
        self._enum_value = KindEnumMember(asdl_name)

    def CStructName(self):
        return self._name.c_name

    def MemberRef(self, member: Member):
        return '{raw_ptr}->{union}.{struct}.{member}'.format(
            raw_ptr=self.GetRawPtr(),
            union=self.UNION_NAME,
            struct=self.CStructName(),
            member=member.CMemberName())

    def HasLocation(self):
        return self._parent.has_location

    def HasEndLocation(self):
        return self._parent.has_end_location

    @property
    def enum_value(self):
        "For `ReturnNode`, this will be `kReturn`"
        return self._enum_value.Name()


class ProductNode(FinalAstNode, Location):
    """
    Things like:
    `comprehension = (expr target, expr iter, expr* ifs, int is_async)`
    """

    MangleName = TypeNode.MangleName

    def __init__(self, asdl_name, members: Sequence[Member],
                 has_location: bool, has_end_location: bool):
        FinalAstNode.__init__(self,
                              name=self.MangleName(asdl_name),
                              parent=AST_NODE,
                              members=members)
        Location.__init__(self, has_location, has_end_location)

    def HasLocation(self):
        return self.has_location

    def HasEndLocation(self):
        return self.has_end_location

    def MemberRef(self, member: Member):
        return '{raw_ptr}->{member}'.format(raw_ptr=self.GetRawPtr(),
                                            member=member.CMemberName())


class AsdlSeq(AstNode):
    """
    `AsdlSeq<[Scalar|Enum]>` mapped to `asdl_int_seq`.
    `AsdlSeq<[Object|AstNode]>` mapped to `asdl_seq`.
    """
    def MangleName(self, asdl_name, element_type):
        return NameBundle(asdl_name + '*',
                          c_name='asdl_int_seq' if isinstance(
                              element_type, (Scalar, Enum)) else 'asdl_seq',
                          cc_name='AsdlSeq<{}>'.format(element_type.name))

    def __init__(self, asdl_name, element_type):
        super().__init__(self.MangleName(asdl_name, element_type),
                         parent=AST_NODE)
        self._element_type = element_type


class EnumMember(Node):
    @abc.abstractclassmethod
    def MangleName(cls, asdl_name):
        "`KindEnumMember` and `SumEnumMember` have different mangling rules"

    def __init__(self, asdl_name):
        super().__init__(self.MangleName(asdl_name))

    def Name(self):
        return self._name.cc_name

    def Value(self):
        return self._name.c_name

    @property
    def value(self):
        return self.Value()


class SumEnumMember(EnumMember):
    @classmethod
    def MangleName(cls, asdl_name):
        "`Add` => `Add` in C."
        return NameBundle(asdl_name, cc_name='k' + asdl_name, c_name=asdl_name)


class KindEnumMember(EnumMember):
    @classmethod
    def MangleName(cls, asdl_name):
        "`Return` => `Return_kind` in C."
        return NameBundle(asdl_name,
                          cc_name='k' + asdl_name,
                          c_name=asdl_name + '_kind')


class Enum(Node):
    @abc.abstractclassmethod
    def MangleName(cls, asdl_name):
        pass

    @abc.abstractclassmethod
    def CreateMember(cls, asdl_name):
        pass

    def __init__(self, asdl_name, member_names: Sequence[str]):
        self._name = self.MangleName(asdl_name)
        self._members = tuple(self.CreateMember(name) for name in member_names)

    @property
    def members(self):
        assert self._members
        return self._members


class KindEnum(Enum):
    @classmethod
    def CreateMember(cls, asdl_name):
        return KindEnumMember(asdl_name)

    @classmethod
    def MangleName(cls, asdl_name):
        """
        `mod` => `_mod_kind` in C.
        => `mod_kind` in Cpp.
        """
        return NameBundle(asdl_name,
                          cc_name=asdl_name + '_kind',
                          c_name='_' + asdl_name + '_kind')


class SumEnum(Enum):
    @classmethod
    def CreateMember(cls, asdl_name):
        return SumEnumMember(asdl_name)

    @classmethod
    def MangleName(cls, asdl_name):
        """
        `operator` => `operator_ty` in C.
        => `operator_kind` in Cpp.
        """
        return NameBundle(asdl_name,
                          cc_name=asdl_name + '_kind',
                          c_name=asdl_name + '_ty')


class Builtin(Node):
    @classmethod
    def MangleName(cls, asdl_name):
        return NameBundle(asdl_name, asdl_name, None)

    def __init__(self, asdl_name: str):
        super().__init__(self.MangleName(asdl_name))

    @classmethod
    def CreateScalar(cls, asdl_name):
        return Scalar(asdl_name)

    @classmethod
    def CreateObject(cls, asdl_name):
        return Object(asdl_name)


class Scalar(Builtin):
    pass


class Object(Builtin):
    def GetRawPtr(self):
        return 'ptr()'


class Hierarchy:
    """
    Contains all the nodes and links in the ast class hierarchy.
    
    Also contains the things that are inconvenient to implement in each of the nodes.
    """
    def __init__(self, type_list):
        self._type_list = type_list

    @property
    def enum_list(self):
        return filter(lambda x: isinstance(x, Enum), self._type_list)

    @property
    def class_list(self):
        return filter(
            lambda x: isinstance(x, AstNode) and not isinstance(x, AsdlSeq),
            self._type_list)

    @property
    def type_list(self):
        return filter(lambda x: isinstance(x, (Enum, AstNode)),
                      self._type_list)


class Config:
    """
    This class contains configuration data.

    Use it as a list:
    `for name in config:`
    `name` will a configuration key that has the implicit value `True`.
    Access to any absent name will result in `False`.
    """

    PREFIX = 'PYAST_HAS'
    HAS_END_LOCATION = PREFIX + '_END_LOCATION'

    @classmethod
    def MangleTypeName(cls, type: Union[Enum, AstNode]):
        """
        `expr` => `PYAST_HAS_expr`
        `Return` => `PYAST_HAS_Return`
        `AsyncFunctionDef` => `PYAST_HAS_AsyncFunctionDef`
        `keyword` => `PYAST_HAS_keyword`
        """
        return '_'.join([cls.PREFIX, type.asdl_name])

    @classmethod
    def MangleMemberName(cls, type: Union[Enum, AstNode],
                         member: Union[EnumMember, Member]):
        """
        `Return.value` => `PYAST_HAS_Return_value`.
        """
        return '_'.join([cls.PREFIX, type.asdl_name, member.asdl_name])

    @classmethod
    def MangleHierarchy(cls, hierarchy: Hierarchy):
        """
        Mangle the whole `hierarchy` of names
        """
        names = []
        for type in hierarchy.type_list:
            typename = cls.MangleTypeName(type)
            names.append(typename)
            if not hasattr(type, 'members'):
                continue
            for m in type.members:
                member_name = cls.MangleMemberName(type, m)
                names.append(member_name)
        return names

    @classmethod
    def Create(cls, hierarchy: Hierarchy):
        # TODO: add HAS_END_LOCATION.
        names = cls.MangleHierarchy(hierarchy)
        return cls(names)

    def __init__(self, names: Sequence[str]):
        self._name_list = tuple(names)
        self._name_set = frozenset(names)

    def __iter__(self):
        return iter(self._name_list)
